package com.haidechizi.tree.code;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * Thread safe
 */
public class HuffmanCode {

    public static Object[] huffmanZip(byte[] bytes) {

        // 统计每个byte出现的次数，并放入一个集合
        List<Node> nodes = createNodes(bytes);
        // 创建一颗哈夫曼数
        Node node = createHuffmanTree(nodes);
        // 创建一个哈夫曼编码表
        Map<Byte, String> huffCodes = createHuffCodes(node);


        // 进行哈夫曼编码
        byte[] b = zip(bytes, huffCodes);

//        write(b, huffCodes, name);

//        byte[] bytes1 = huffmanDecode(b, huffCodes);

//        String ss = new String(bytes1);

        return new Object[]{b, huffCodes};
    }

    private static byte[] zip(byte[] bytes, Map<Byte, String> huffCodes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(huffCodes.get(b));
        }
        System.out.println(sb.toString());
        // 将二进制字符串转为byte
        int length = sb.length() / 8;
        if (sb.length() % 8 != 0) {
            length += 1;
        }
        byte[] by = new byte[length];
        int index = 0;
        for (int i = 0; i < sb.length(); i += 8) {
            String byteStr;
            if (i + 8 > sb.length()) {
                byteStr = sb.substring(i);
            } else {
                byteStr = sb.substring(i, i + 8);
            }


            byte b = (byte) Integer.parseInt(byteStr, 2);
            by[index] = b;
            index++;
        }
        return by;
    }

    private static Map<Byte, String> createHuffCodes(Node node) {
        StringBuilder sb = new StringBuilder();
        Map<Byte, String> huffCodes = new HashMap<>();
        createCodes(node.getLeft(), "0", sb, huffCodes);
        createCodes(node.getRight(), "1", sb, huffCodes);
        return huffCodes;
    }

    private static void createCodes(Node node, String s, StringBuilder sb, Map<Byte, String> huffCodes) {
        if (node == null) {
            return;
        }
        StringBuilder stringBuilder = new StringBuilder(sb);
        stringBuilder.append(s);
        // 非叶子节点
        if (node.getData() == null) {
            createCodes(node.getLeft(), "0", stringBuilder, huffCodes);
            createCodes(node.getRight(), "1", stringBuilder, huffCodes);
        } else {
            huffCodes.put(node.getData(), stringBuilder.toString());
        }
    }

    private static Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            // 对集合进行排序（倒序）
            Collections.sort(nodes);
            // 取出最后连个节点，并移除
            Node left = nodes.remove(nodes.size() - 1);
            Node right = nodes.remove(nodes.size() - 1);
            //Node left = nodes.get(nodes.size() - 1);
            //Node right = nodes.get(nodes.size() - 2);
            Node parent = new Node(null, left.getWeight() + right.getWeight());

            //nodes.remove(left);
            // nodes.remove(right);
            parent.setLeft(left);
            parent.setRight(right);
            // 将新的节点放入集合中
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    private static List<Node> createNodes(byte[] bytes) {
        List<Node> nodes = new ArrayList<>();
        // 统计每个byte出现的次数
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null) {
                counts.put(b, 1);
            } else {
                counts.put(b, count + 1);
            }
        }
        // 构建node并放入到集合中
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }


    public static byte[] huffmanDecode(byte[] bytes, Map<Byte, String> huffcodes) {
        // 将byte数组转为二进制字符串
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < bytes.length; i++) {
            boolean flag = (i == bytes.length - 1);
            String str = byteToBitStr(bytes[i], flag);
            sb.append(str);
        }
        System.out.println(sb.toString());

        // 将哈夫曼编码表调换位置
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffcodes.entrySet()) {
            map.put(entry.getValue(), entry.getKey());
        }

        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < sb.length(); ) {
            int count = 1;
            boolean flag = true;
            while (flag) {
                String str = sb.substring(i, i + count);
                Byte aByte = map.get(str);
                if (aByte != null) {
                    flag = false;
                    list.add(aByte);
                } else {
                    count++;
                }
            }
            i += count;
        }
        byte[] byts = new byte[list.size()];
        for (int i = 0; i < byts.length; i++) {
            byts[i] = list.get(i);
        }
        return byts;
    }

    private static String byteToBitStr(byte b, boolean flag) {
        int temp = b;
        if (!flag) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str;
        } else {
            return str.substring(str.length() - 8);

        }

    }

    public static void zipFile(String name, String outPath) {
        File file = new File(name);
        try (FileInputStream fis = new FileInputStream(file);
             FileOutputStream fos = new FileOutputStream(new File(outPath));
             ObjectOutputStream oos = new ObjectOutputStream(fos)) {
            byte[] bytes = new byte[fis.available()];
            fis.read(bytes);
            Object[] obj = huffmanZip(bytes);
            String fileName = file.getName();
            oos.writeObject(obj[0]);
            oos.writeObject(obj[1]);
            oos.writeObject(fileName);
        } catch (Exception e) {

        }
    }

    public static void unzipFile(String filePath, String outPath) {
        File file = new File(filePath);
        try (FileInputStream fis = new FileInputStream(file);
             ObjectInputStream ois = new ObjectInputStream(fis)
        ) {
            byte[] datas = (byte[]) ois.readObject();
            Map<Byte, String> map = (Map<Byte, String>) ois.readObject();
            String fileName = (String) ois.readObject();
            byte[] bytes = huffmanDecode(datas, map);
            FileOutputStream fos = new FileOutputStream(new File(outPath + "/" + fileName));
            ObjectOutputStream oos = new ObjectOutputStream(fos);
            oos.writeObject(bytes);
            oos.flush();
            oos.close();
            fos.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


}
